import torch.nn as nn
from model.encoder import Encoder
from model.decoder import Decoder
from model.attack import Attack
from options import HiDDenConfiguration
from noise_layers.noiser import Noiser
from noise_layers.gaussian_filter import GF
from noise_layers.gaussian_noise import Gaussian_Noise
from noise_layers.colorjitter import ColorJitter
from noise_layers.salt_pepper_noise import SP
from noise_layers.jpeg import Jpeg, JpegSS, JpegMask
from noise_layers.quantization import Quantization
from noise_layers.jpeg_compression import JpegCompression
from modules.UnetPlusPlus_L4 import UNetPlusPlusL4_Attack
from modules.UnetPlusPlus_L5 import UNetPlusPlusL5_Attack
from modules.UnetPlusPlus_L6 import UNetPlusPlusL6_Attack
from modules.UnetAttack import UNetAttack
from modules.UNetPlus3.UNet3Plus import UNet3Plus
from modules.TransUNet import TransUNet_Attack
from modules.ParallelImageNet import ParallelImageNet
from HiNet.model import Model, init_model

class EncoderDecoder(nn.Module):
    """
    Combines Encoder->Noiser->Decoder into single pipeline.
    The input is the cover image and the watermark message. The module inserts the watermark into the image
    (obtaining encoded_image), then applies Noise layers (obtaining noised_image), then passes the noised_image
    to the Decoder which tries to recover the watermark (called decoded_message). The module outputs
    a three-tuple: (encoded_image, noised_image, decoded_message)
    """
    def __init__(self):

        super(EncoderDecoder, self).__init__()
        self.model = Model().cuda()

    def forward(self, image, message):
        pass